Summary

• Adds a with_sharding_constraint operator similar to JAX's API for constraining intermediate tensor shardings
• Uses local_map internally with an identity function for a simple implementation
• Exports the new operator from the public API

Implementation Details

The with_sharding_constraint(x, shardings, device_mesh=None) function:

• Takes a tensor, desired placements, and optional device mesh
• Uses local_map with matching in_placements and out_placements to enforce the constraint
• Falls back to global mesh from enclosing local_map region if device_mesh is not provided
• Uses clone() inside the identity function to avoid input-to-output aliasing issues with dynamo's HOP tracing

Example Usage

from autoparallel import with_sharding_constraint
from torch.distributed.tensor.placement_types import Shard, Replicate

# Constrain intermediate tensor to be sharded
x = with_sharding_constraint(x, (Shard(0), Replicate()), device_mesh)

Test plan

• Test with explicit device mesh
• Test between local_map regions
• Test forcing replication
• Test with 2D mesh
• Test multiple constraints in sequence
• Test error when no mesh is available
• Verify placements match requested constraints in optimizer output